import time
import json
import random
import websocket
from locust import HttpUser, task, between, events

class WebSocketClient:
    def __init__(self, host):
        self.host = host
        self.ws = None
        self.connected = False
        self.user_id = None
        self.token = None
        self.character_id = None
    
    def connect(self, user_id=None, token=None):
        """建立WebSocket连接"""
        try:
            # 如果有登录凭证，则在URL中添加
            url = f"ws://{self.host}/game/socket"
            if user_id and token:
                url += f"?user_id={user_id}&token={token}"
            
            self.ws = websocket.create_connection(url)
            self.connected = True
            self.user_id = user_id
            self.token = token
            return True
        except Exception as e:
            print(f"WebSocket连接失败: {e}")
            return False
    
    def disconnect(self):
        """关闭WebSocket连接"""
        if self.ws and self.connected:
            self.ws.close()
            self.connected = False
    
    def send_message(self, message_type, data):
        """发送WebSocket消息"""
        if not self.connected:
            print("WebSocket未连接，无法发送消息")
            return False
        
        message = {
            "type": message_type,
            "data": data,
            "timestamp": int(time.time() * 1000)
        }
        
        try:
            self.ws.send(json.dumps(message))
            return True
        except Exception as e:
            print(f"发送消息失败: {e}")
            return False
    
    def receive_message(self, timeout=5):
        """接收WebSocket消息"""
        if not self.connected:
            print("WebSocket未连接，无法接收消息")
            return None
        
        try:
            self.ws.settimeout(timeout)
            response = self.ws.recv()
            return json.loads(response)
        except Exception as e:
            print(f"接收消息失败: {e}")
            return None


class GameUser(HttpUser):
    wait_time = between(1, 3)  # 用户操作之间的等待时间
    
    def on_start(self):
        """用户会话开始时执行"""
        self.ws_client = WebSocketClient(self.host.replace("http://", "").replace("https://", ""))
        self.username = f"test_user_{random.randint(1000, 9999)}"
        self.password = "password123"
        self.logged_in = False
        self.in_game = False
        
        # 注册新用户（如果需要）
        self.register()
        
        # 登录
        self.login()
        
        # 连接WebSocket
        if self.logged_in:
            self.connect_websocket()
            
            # 创建角色（如果需要）
            if not self.character_id:
                self.create_character()
            
            # 选择角色进入游戏
            if self.character_id:
                self.enter_game()
    
    def on_stop(self):
        """用户会话结束时执行"""
        # 登出游戏
        if self.in_game:
            self.logout_game()
        
        # 断开WebSocket连接
        if self.ws_client.connected:
            self.ws_client.disconnect()
    
    def register(self):
        """注册新用户"""
        try:
            with self.client.post("/api/register", json={
                "username": self.username,
                "password": self.password,
                "email": f"{self.username}@example.com"
            }, catch_response=True) as response:
                if response.status_code == 200:
                    result = response.json()
                    if result.get("success"):
                        print(f"用户 {self.username} 注册成功")
                    else:
                        print(f"用户注册失败: {result.get('message', '未知错误')}")
                else:
                    print(f"用户注册请求失败: HTTP {response.status_code}")
        except Exception as e:
            print(f"注册过程中发生错误: {e}")
    
    def login(self):
        """用户登录"""
        try:
            with self.client.post("/api/login", json={
                "username": self.username,
                "password": self.password
            }, catch_response=True) as response:
                if response.status_code == 200:
                    result = response.json()
                    if result.get("success"):
                        self.user_id = result.get("user_id")
                        self.token = result.get("token")
                        self.logged_in = True
                        print(f"用户 {self.username} 登录成功")
                    else:
                        print(f"用户登录失败: {result.get('message', '未知错误')}")
                else:
                    print(f"用户登录请求失败: HTTP {response.status_code}")
        except Exception as e:
            print(f"登录过程中发生错误: {e}")
    
    def connect_websocket(self):
        """连接WebSocket"""
        if self.ws_client.connect(self.user_id, self.token):
            print(f"用户 {self.username} WebSocket连接成功")
        else:
            print(f"用户 {self.username} WebSocket连接失败")
    
    def create_character(self):
        """创建游戏角色"""
        character_name = f"hero_{random.randint(1000, 9999)}"
        character_class = random.choice(["warrior", "mage", "archer", "priest"])
        
        success = self.ws_client.send_message("create_character", {
            "name": character_name,
            "class": character_class,
            "appearance": {
                "hair": random.randint(1, 10),
                "face": random.randint(1, 10),
                "color": random.randint(1, 10)
            }
        })
        
        if success:
            # 等待服务器响应
            response = self.ws_client.receive_message()
            if response and response.get("type") == "character_created" and response.get("success"):
                self.character_id = response.get("data", {}).get("character_id")
                print(f"角色 {character_name} 创建成功，ID: {self.character_id}")
            else:
                print(f"角色创建失败: {response.get('message', '未知错误') if response else '无响应'}")
    
    def enter_game(self):
        """选择角色进入游戏"""
        if not self.character_id:
            print("没有可用的角色，无法进入游戏")
            return
        
        success = self.ws_client.send_message("enter_game", {
            "character_id": self.character_id
        })
        
        if success:
            # 等待服务器响应
            response = self.ws_client.receive_message()
            if response and response.get("type") == "game_entered" and response.get("success"):
                self.in_game = True
                print(f"角色 {self.character_id} 成功进入游戏")
            else:
                print(f"进入游戏失败: {response.get('message', '未知错误') if response else '无响应'}")
    
    def logout_game(self):
        """退出游戏"""
        success = self.ws_client.send_message("logout", {})
        if success:
            self.in_game = False
            print(f"角色 {self.character_id} 成功退出游戏")
    
    @task(3)
    def move_character(self):
        """移动角色"""
        if not self.in_game:
            return
        
        # 随机移动方向和距离
        direction = random.choice(["north", "south", "east", "west"])
        distance = random.randint(1, 10)
        
        success = self.ws_client.send_message("move", {
            "direction": direction,
            "distance": distance
        })
        
        if success:
            response = self.ws_client.receive_message(timeout=2)
            if response and response.get("type") == "move_result" and response.get("success"):
                new_position = response.get("data", {}).get("position", {})
                print(f"角色移动成功，新位置: X={new_position.get('x')}, Y={new_position.get('y')}")
            else:
                if response:
                    print(f"角色移动失败: {response.get('message', '未知错误')}")
    
    @task(1)
    def chat_message(self):
        """发送聊天消息"""
        if not self.in_game:
            return
        
        messages = [
            "你好！", "有人组队吗？", "这个游戏真好玩！", 
            "请问新手任务在哪里？", "有人知道怎么升级快吗？"
        ]
        
        success = self.ws_client.send_message("chat", {
            "channel": "world",
            "message": random.choice(messages)
        })
        
        if success:
            print(f"聊天消息发送成功")
    
    @task(2)
    def check_surroundings(self):
        """查看周围环境"""
        if not self.in_game:
            return
        
        success = self.ws_client.send_message("get_surroundings", {
            "range": random.randint(5, 20)
        })
        
        if success:
            response = self.ws_client.receive_message(timeout=2)
            if response and response.get("type") == "surroundings" and response.get("success"):
                entities = response.get("data", {}).get("entities", [])
                print(f"周围有 {len(entities)} 个实体")
            else:
                if response:
                    print(f"获取周围环境失败: {response.get('message', '未知错误')}")


# 自定义事件处理
@events.init.add_listener
def on_locust_init(environment, **kwargs):
    print("Locust初始化完成，准备开始测试...")

@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    print("测试开始...")

@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    print("测试结束...")


if __name__ == "__main__":
    # 本地测试用
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "debug":
        # 创建一个测试用户并执行完整流程
        user = GameUser(None)
        user.host = "localhost:8080"  # 替换为实际的游戏服务器地址
        user.on_start()
        
        # 执行几次任务
        for _ in range(5):
            user.move_character()
            time.sleep(1)
            user.chat_message()
            time.sleep(1)
            user.check_surroundings()
            time.sleep(1)
        
        user.on_stop()